package com.my.collection.tree.v3;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

/**
 * Created by ZhongGang
 * at 2018/10/29 11:36
 */
public class BinarySearchTree<E> {
    private Node root;
    private Comparator<E> comparator;

    public BinarySearchTree() {
        this.comparator = new Comparator<E>() {
            @Override
            public int compare(E o1, E o2) {
                Comparable<E> c1 = (Comparable<E>) o1;
                return c1.compareTo(o2);
            }
        };
    }

    public BinarySearchTree(Node root) {
        this();
        this.root = root;
    }

    public BinarySearchTree(Node root, Comparator<E> comparator) {
        this.root = root;
        this.comparator = comparator;
    }

    public boolean add(E e) {
        return add(e, this.root);
    }

    private boolean add(E e, Node node) {
        if (node == null) {
            this.root = new Node(e, null, null);
            return true;
        }

        if (this.comparator.compare(e, node.data) < 0) {
            if (node.left != null) {
                add(e, node.left);
            } else {
                node.left = new Node(e, null, null);
                return true;
            }
        } else if (this.comparator.compare(e, node.data) > 0) {
            if (node.right != null) {
                add(e, node.right);
            } else {
                node.right = new Node(e, null, null);
                return true;
            }
        } else {
            return false;
        }

        return false;
    }

    public boolean contain(E e) {
        return contain(e, this.root);
    }

    private boolean contain(E e, Node node) {
        if (node == null) {
            return false;
        }

        if (this.comparator.compare(e, node.data) < 0) {
            return contain(e, node.left);
        } else if (this.comparator.compare(e, node.data) > 0) {
            return contain(e, node.right);
        } else {
            return true;
        }
    }

    public boolean remove(E e) {
        throw new UnsupportedOperationException();
    }

    public E min() {
        return min(this.root);
    }

    private E min(Node node) {
        if (node == null) {
            return null;
        }

        if (node.left == null) {
            return node.data;
        }

        return min(node.left);
    }

    public E max() {
        return max(this.root);
    }

    private E max(Node node) {
        if (node == null) {
            return null;
        }

        if (node.right == null) {
            return node.data;
        }

        return max(node.right);
    }

    public E floor(E e) {
        return floor(e, this.root);
    }

    private E floor(E e, Node node) {
        if (node == null) {
            return null;
        }

        if (this.comparator.compare(e, node.data) < 0) {
            if (node.left != null) {
                if (this.comparator.compare(e, node.left.data) > 0) {
                    return node.left.data;
                } else {
                    return floor(e, node.left);
                }
            } else {
                return null;
            }
        } else if (this.comparator.compare(e, node.data) > 0) {
            if (node.right != null) {
                if (this.comparator.compare(e, node.right.data) < 0) {
                    return node.data;
                } else {
                    return floor(e, node.right);
                }
            } else {
                return node.data;
            }
        } else {
            return e;
        }
    }

    public E ceil(E e) {
        return ceil(e, this.root);
    }

    private E ceil(E e, Node node) {
        if (node == null) {
            return null;
        }

        if (this.comparator.compare(e, node.data) < 0) {
            if (node.left != null) {
                if (this.comparator.compare(e, node.left.data) > 0) {
                    return node.data;
                } else {
                    return ceil(e, node.left);
                }
            } else {
                return node.data;
            }
        } else if (this.comparator.compare(e, node.data) > 0) {
            if (node.right != null) {
                if (this.comparator.compare(e, node.right.data) < 0) {
                    return node.right.data;
                } else {
                    return ceil(e, node.right);
                }
            } else {
                return null;
            }
        } else {
            return e;
        }
    }

    public List<E> less(E e) {
        List<E> results = new ArrayList<>();
        less(e, this.root, results);
        return results;
    }

    private void less(E e, Node node, List<E> results) {
        if (node == null) {
            return;
        }

        int compare = this.comparator.compare(e, node.data);
        if (compare < 0) {
            less(e, node.left, results);
        } else if (compare > 0) {
            results.add(node.data);
            less(e, node.left, results);
            less(e, node.right, results);
        } else {
            less(e, node.left, results);
        }
    }

    public List<E> greater(E e) {
        List<E> results = new ArrayList<>();
        greater(e, this.root, results);
        return results;
    }

    private void greater(E e, Node node, List<E> results) {
        if (node == null) {
            return;
        }

        int compare = this.comparator.compare(e, node.data);
        if (compare < 0) {
            results.add(node.data);
            greater(e, node.left, results);
            greater(e, node.right, results);
        } else if (compare > 0) {
            greater(e, node.right, results);
        } else {
            greater(e, node.right, results);
        }
    }

    private class Node implements Comparable<Node> {
        E data;
        Node left;
        Node right;

        public Node(E data, Node left, Node right) {
            this.data = data;
            this.left = left;
            this.right = right;
        }

        @Override
        public int compareTo(Node o) {
            return BinarySearchTree.this.comparator.compare(this.data, o.data);
        }
    }

}
